"""
des: 
author: 邵宁
date: 2022-09-10
"""
import os
import pickle
from util import db, path


def msg_read(msg_token):
    msg_path = "static/session/" + msg_token
    msg_path = os.path.join(msg_path, "msg_obj")
    if not os.path.exists(msg_path):
        return []
    with open(msg_path, 'rb') as f:
        return pickle.load(f)


def msg_write(msg, msg_token):
    if msg_token == "":
        return
    msg_pre = msg_read(msg_token)
    msg_pre.append(msg)
    msg_path = "static/session/" + msg_token
    msg_path_obj = os.path.join(msg_path, "msg_obj")
    path.dir_init(msg_path_obj)
    with open(msg_path_obj, 'wb') as f:
        pickle.dump(msg_pre, f)
    db.update_session_msg_tag(msg_token)


def setting_read(session_token):
    setting_path = "static/session/" + session_token
    setting_path = os.path.join(setting_path, "setting_obj")
    if not os.path.exists(setting_path):
        return {}
    with open(setting_path, 'rb') as f:
        return pickle.load(f)


def setting_write(setting, session_token):
    setting_pre = setting_read(session_token)
    setting_pre.update(setting)
    if setting_pre.get("codeName", "") != "":
        db.update_session_name(session_token, setting_pre.get("codeName", ""))
    setting_path = "static/session/" + session_token
    setting_path_obj = os.path.join(setting_path, "setting_obj")
    path.dir_init(setting_path_obj)
    with open(setting_path_obj, 'wb') as f:
        pickle.dump(setting_pre, f)


class MsgQueue:
    def __init__(self):
        self.msg_queue = {}  # task_id: []

    def push(self, task_id, msg_obj):
        if task_id in self.msg_queue:
            self.msg_queue[task_id].append(msg_obj)
        else:
            self.msg_queue[task_id] = []

    def pop(self, task_id):
        if task_id in self.msg_queue and len(self.msg_queue[task_id]) > 0:
            msg_obj = self.msg_queue[task_id].pop(0)
            if msg_obj.get("status_is_end"):
                self.msg_queue.pop(task_id)
            return msg_obj
        return {}

    def pop_batch(self, task_id):
        msg_obj_list = []
        if task_id in self.msg_queue:
            for _ in range(len(self.msg_queue[task_id])):
                msg_obj_list.append(self.msg_queue[task_id].pop(0))
        return msg_obj_list

    def push_msg(self, task_id, msg: str, is_end=False):
        msg_obj = {
            "msg_id": len(self.msg_queue[task_id]) + 1 if task_id in self.msg_queue else 0,
            "msg_type": "msg",
            "timestamp": path.get_time_stamp(),
            "status_is_end": is_end,
            "msg": msg
        }
        self.push(task_id, msg_obj)

    def push_link(self, task_id, link: str, filename="未命名", is_end=False):
        msg_obj = {
            "msg_id": len(self.msg_queue[task_id]) + 1 if task_id in self.msg_queue else 0,
            "msg_type": "link",
            "timestamp": path.get_time_stamp(),
            "status_is_end": is_end,
            "msg": "",
            "link": link,
            "filename": filename,
        }
        self.push(task_id, msg_obj)
